import numpy as np
import sympy as sp
import sys

sys.path.append('../pde')
sys.path.append('../fem')

from fealpy.mesh import MeshFactory
from Linearelasticitymodelpde import PDE
from fealpy.pde.linear_elasticity_model import QiModel3d, PolyModel3d, Model2d, HuangModel2d
from LinearElasticityFEMModel import LinearElasticityFEMModel
from fealpy.tools.show import showmultirate
from timeit import default_timer as timer

import matlab.engine
import matplotlib.pyplot as plt

p = int(sys.argv[1])
n = int(sys.argv[2])
maxit = int(sys.argv[3])

pi = sp.pi
sin = sp.sin
cos = sp.cos
exp = sp.exp
ln = sp.ln
mf = MeshFactory()
start = timer()
eng = matlab.engine.start_matlab()
end = timer()
print('Start matlab engine time:',end-start)
eng.addpath('/Users/chen/git/mixed-hzelem/test/solver')



################################
###########由真解创建pde########## 
x = sp.symbols('x0:3')
#u = [exp(x[0]+x[1]),exp(x[0]-x[1])]
u = [exp(x[0]+x[1]),exp(x[0]-x[1]),x[1]+x[2]]
#u = [sin(2*pi*x[0]),sin(2*pi*x[1]),sin(2*pi*x[2])]
#u = [sin(2*pi*x[0])*sin(2*pi*x[1])*sin(2*pi*x[2]),sin(2*pi*x[0])*sin(2*pi*x[1])*sin(2*pi*x[2]),sin(2*pi*x[0])*sin(2*pi*x[1])*sin(2*pi*x[2]),]

PDEs = PDE(u,x,lam=10)
pde = PDEs.Displacement_bc()
#pde = PDEs.Stress_and_Displacement_bc()
gdim = PDEs.gdim

################################
###########加载mesh#############
if gdim == 2:
    mesh = mf.boxmesh2d(pde.domain(),nx=1,ny=1,meshtype='tri')
elif gdim == 3:
    mesh = mf.boxmesh3d(pde.domain(),nx=1,ny=1,nz=1,meshtype='tet')


mesh.uniform_refine(n)
integrator = mesh.integrator(5)

###################################
##############fem求解##############
errorType = ['$||\sigma - \sigma_h ||_{0}$',
                '$||div(\sigma - \sigma_h)||_{0}$',
                '$||u - u_h||_{0}$',
                '$||\sigma - \sigma_I ||_{0}$',
                '$||div(\sigma - \sigma_I)||_{0}$',
                '$||u - u_I||_{0}$',
                '$|| u_I - u_h||_{0}$'
                ]
Ndof = np.zeros((maxit,))
NC = np.zeros((maxit,))
errorMatrix = np.zeros((len(errorType), maxit), dtype=np.float)



for i in range(maxit):
    print(i+1)
    fem = LinearElasticityFEMModel(mesh, pde, p, integrator)
    fem.get_AA_and_bb()
    fem.Save_A_and_b_to_matlab()
    fem.precondieitoner()
    fem.Save_A_and_b_to_matlab_fast()
    #x = np.array(eng.matlab_solve())[:,0];
    x = np.array(eng.matlab_fast_solve())[:,0]
    fem.Finite_emelent_solution(x=x)
    Ndof[i] = len(x)
    NC[i] = mesh.number_of_cells()
    errorMatrix[:, i] = fem.error()
    if i < maxit - 1:
            mesh.uniform_refine()

print('Ndof:', Ndof)
print('Ncell:',NC)
print('error:', errorMatrix)
showmultirate(plt, 0, Ndof, errorMatrix, errorType)
plt.show()

    










